import pulp
import pandas as pd

file_path = '附件1.xlsx'
df_land = pd.read_excel(file_path, sheet_name='乡村的现有耕地')

file_path_2 = '附件2.xlsx'
df_crops = pd.read_excel(file_path_2, sheet_name='2023年统计的相关数据')


# 定义豆类作物的编号集合
bean_crops = {1, 2, 3, 4, 5}
bean_vegs = {17, 18, 19}

A = {row['地块名称']: row['地块面积/亩'] for _, row in df_land.iterrows()}


plots_single_season = df_land[df_land['地块类型'].isin(['平旱地', '梯田', '山坡地'])]['地块名称'].tolist()

plots_double_season = df_land[df_land['地块类型'].isin(['水浇地', '普通大棚', '智慧大棚'])]['地块名称'].tolist()

irrigated_land = df_land[df_land['地块类型'] == '水浇地']['地块名称'].tolist()
greenhouse_land = df_land[df_land['地块类型'] == '普通大棚']['地块名称'].tolist()
smart_greenhouse_land = df_land[df_land['地块类型'] == '智慧大棚']['地块名称'].tolist()

model = pulp.LpProblem("Crop_Planting_Optimization", pulp.LpMaximize)

crops_single_season = list(range(1, 16))  # 作物1-15，适用于单季地块
crops_double_season_irrigated = list(range(16, 38))  # 作物16-37，适用于水浇地
crops_double_season_greenhouse = list(range(23, 42))  # 作物23-41，适用于普通大棚
crops_double_season_smart = list(range(27, 35))  # 作物27-34，适用于智慧大棚

df_crops['销售单价(元/斤)'] = df_crops['销售单价/(元/斤)'].apply(
    lambda x: (float(x.split('-')[0]) + float(x.split('-')[1])) / 2 if isinstance(x, str) else x
)

P = {row['作物编号']: row['销售单价(元/斤)'] for _, row in df_crops.iterrows()}  # 销售单价字典
Y = {row['作物编号']: row['亩产量/斤'] for _, row in df_crops.iterrows()}  # 亩产量字典
C = {row['作物编号']: row['种植成本/(元/亩)'] for _, row in df_crops.iterrows()}  # 种植成本字典

years = list(range(2024, 2031))  # 从2024到2030
x = pulp.LpVariable.dicts("x", (plots_single_season + plots_double_season,
                                crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart,
                                years, [1, 2]), 0, 1, cat='Continuous')
y1 = pulp.LpVariable.dicts("y1", (plots_single_season + plots_double_season,
                                  crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart,
                                  years, [1, 2]), 0, 1, cat='Binary')

y2 = pulp.LpVariable.dicts("y2", (plots_single_season + plots_double_season,
                                  crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart,
                                  years, [1, 2]), 0, 1, cat='Binary')

# 目标函数：最大化所有地块的净收益
# 单季地块目标函数
# 单季地块目标函数
Z_single = pulp.lpSum(
    (P[k] * Y[k] * A[i] - C[k] * A[i]) * x[i][k][t][1]
    for t in years
    for i in plots_single_season
    for k in crops_single_season
)

# 双季地块目标函数
Z_irrigated = pulp.lpSum(
    (P[k] * Y[k] * A[i] - C[k] * A[i]) * (x[i][k][t][1] + x[i][k][t][2])
    for t in years
    for i in irrigated_land
    for k in crops_double_season_irrigated
)

Z_greenhouse = pulp.lpSum(
    (P[k] * Y[k] * A[i] - C[k] * A[i]) * (x[i][k][t][1] + x[i][k][t][2])
    for t in years
    for i in greenhouse_land
    for k in crops_double_season_greenhouse
)

Z_smart_greenhouse = pulp.lpSum(
    (P[k] * Y[k] * A[i] - C[k] * A[i]) * (x[i][k][t][1] + x[i][k][t][2])
    for t in years
    for i in smart_greenhouse_land
    for k in crops_double_season_smart
)

model += Z_single + Z_irrigated + Z_greenhouse + Z_smart_greenhouse
for i in plots_single_season + irrigated_land + greenhouse_land + smart_greenhouse_land:
    for k in crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart:
        for t in years:
            for j in [1, 2]:

                model += x[i][k][t][j] == 0.5 * y1[i][k][t][j] + y2[i][k][t][j]
                model += y1[i][k][t][j] + y2[i][k][t][j] <= 1
for t in years:
    for i in plots_single_season:
        model += pulp.lpSum(x[i][k][t][1] for k in crops_single_season) == 1
    for i in irrigated_land:
        model += pulp.lpSum(x[i][k][t][1] for k in range(16, 35)) == 1
        model += pulp.lpSum(x[i][k][t][2] for k in [16, 35, 36, 37]) == 1

    for i in greenhouse_land:
        model += pulp.lpSum(x[i][k][t][1] for k in range(17, 35)) == 1
        model += pulp.lpSum(x[i][k][t][2] for k in range(38, 42)) == 1

    for i in smart_greenhouse_land:
        model += pulp.lpSum(x[i][k][t][1] for k in range(17, 35)) == 1
        model += pulp.lpSum(x[i][k][t][2] for k in range(17, 35)) == 1

for t in years:
    for i in plots_single_season:
        model += pulp.lpSum(x[i][k][t][1] for k in crops_single_season) == 1

    for i in plots_double_season:
        for j in [1, 2]:
            model += pulp.lpSum(x[i][k][t][j] for k in crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart) == 1

for t in years[:-1]:
    for i in plots_single_season:
        for k in crops_single_season:
            model += x[i][k][t][1] + x[i][k][t+1][1] <= 1

    for i in plots_double_season:
        for k in crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart:
            model += x[i][k][t][1] + x[i][k][t][2] <= 1
            model += x[i][k][t][2] + x[i][k][t+1][1] <= 1


# 轮作约束：每块地在每连续三年内至少种植一次豆类作物
for t_start in range(2024, 2028):
    for i in plots_single_season:
        model += pulp.lpSum(x[i][k][t][1] for k in bean_crops for t in range(t_start, t_start + 3)) >= 1

    for i in irrigated_land + greenhouse_land + smart_greenhouse_land:
        model += pulp.lpSum(x[i][k][t][1] + x[i][k][t][2] for k in bean_vegs for t in range(t_start, t_start + 3)) >= 1

model.solve()



results = []

for i in plots_single_season + plots_double_season:
    for k in crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart:
        for t in years:
            for j in [1, 2]:
                if x[i][k][t][j].varValue is not None and x[i][k][t][j].varValue > 0:
                    results.append([i, k, t, j, x[i][k][t][j].varValue])

df_results = pd.DataFrame(results, columns=["地块", "作物", "年份", "季节", "种植比例"])

# print(df_results)


# 读取附件数据
df_land = pd.read_excel('附件1.xlsx', sheet_name='乡村的现有耕地')
df_crops = pd.read_excel('附件2.xlsx', sheet_name='2023年统计的相关数据')
df_crop_names = pd.read_excel('附件1.xlsx', sheet_name='乡村种植的农作物')


result = pd.read_excel('result3.xlsx', sheet_name=None)

# 读取作物名称与编号的对应关系
crop_name_mapping = {row['作物编号']: row['作物名称'] for _, row in df_crop_names.iterrows()}
place_space_mapping = {row['地块名称']: row['地块面积/亩'] for _, row in df_land.iterrows()}

print(type(df_results))
# 遍历每年的优化结果并填入 result2.xlsx 对应的表格
for year in range(2024, 2031):
    # 读取该年份对应的表格
    sheet_name = f'{year}'
    sheet_data = result[sheet_name]
    # index0 = sheet_data.index[sheet_data['地块名'] == 'A1'].tolist()
    # print(index0, type(index0))

    # 遍历 df_results 中对应年份的数据
    for _, row in df_results[df_results['年份'] == year].iterrows():
        plot_name = row['地块']  # 地块名称
        crop_id = row['作物']  # 作物编号
        season = row['季节']   # 第几季种植
        planting_ratio = row['种植比例']  # 种植比例
        # print(planting_ratio, type(planting_ratio))
        # 获取作物名称
        crop_name = crop_name_mapping.get(crop_id, None)
        if crop_name is None:
            continue
        land_space = place_space_mapping.get(plot_name, 0)
        if land_space == 0:
            continue
        index_list = sheet_data.index[sheet_data['地块名']==plot_name].tolist()
        index_row = index_list[season-1]
        # 根据季节和地块名称直接插入种植比例
        sheet_data.loc[index_row, crop_name] = planting_ratio * land_space

    # 更新 result2 中该年的数据
    result[sheet_name] = sheet_data

# 保存修改后的 result2.xlsx
with pd.ExcelWriter('result3.xlsx', engine='openpyxl') as writer:
    for sheet_name, data in result.items():
        data.to_excel(writer, sheet_name=sheet_name, index=False)

print("已将结果插入到 result")